from __future__ import division

import sys, re, argparse, logging, csv

from numpy import NaN, Infinity

import common

logger = logging.getLogger("cycle_filter")
logger.addHandler(logging.StreamHandler())
logger.setLevel(logging.INFO)

path_pattern = re.compile(r'<element.* path="(?P<path>[^"]*).*')


class IterationSummary(object):

    def __init__(self, last_peak_extension_m, max_extension_m, distance_n):
        common.check_False(last_peak_extension_m > max_extension_m)
        self.last_peak_extension_m = last_peak_extension_m
        common.check_ge(max_extension_m, 0)
        self.max_extension_m = max_extension_m
        common.check_ge(distance_n, 0)
        self.distance_n = distance_n


class CycleSummary(object):

    def __init__(self, path, iterations):
        self.path = path
        common.check_ge(len(iterations), 1)
        self.iterations = iterations


def main(params):
    """
    Single-iteration curves are ok if they have at least one peak.
    """
    parser = argparse.ArgumentParser()
    parser.add_argument("--cycle_summary_path")
    parser.add_argument("--max_extension_to_end_nm", type=float)
    parser.add_argument("--known_curves_hooke_playlist")
    parser.add_argument("--distance_ordered_csv")
    parser.add_argument("--distance_ordered_tsv")

    args = parser.parse_args(params)
    print(args)

    with open(args.known_curves_hooke_playlist, "rb") as in_file:
        known_good = []
        for line in in_file:
            matcher = path_pattern.match(line)
            if matcher:
                known_good.append(matcher.group("path"))
    l = len(known_good)
    logger.info("len(known_good)=%d", l)
    known_good = set(known_good)
    common.check_equal(len(known_good), l)
    
    with open(args.cycle_summary_path, "rb") as in_file:
        cycle_summaries = [eval(line) for line in in_file]

    all_paths = [cycle_summary.path for cycle_summary in cycle_summaries]
    l = len(all_paths)
    all_paths = set(all_paths)
    common.check_equal(len(all_paths), l)

    diff = known_good - all_paths
    if diff:
        sorted_diff = sorted(diff)
        logger.warning("%d known good without cycle summaries. In sorted order, the first one is %s and the last one is %s.", len(diff), sorted_diff[0], sorted_diff[-1])
    
    max_extension_to_end_m = args.max_extension_to_end_nm / 1e9
    oks = []
    for cycle_summary in cycle_summaries:
        max_extension_except_last_m = max(
            [0] + [
                iteration.max_extension_m
                for iteration in cycle_summary.iterations[:-1]])
        a = cycle_summary.iterations[-1].last_peak_extension_m >= max_extension_except_last_m
        b = all([
                iteration.last_peak_extension_m + max_extension_to_end_m >= iteration.max_extension_m
                for iteration in cycle_summary.iterations[:-1]])
        oks.append(a and b)

    cycle_summary_ok_list = list(zip(cycle_summaries, oks))
    
    cycle_summary_ok_max_distance_list = [
        (cycle_summary, ok, max(
                [0] + [
                    iteration.distance_n
                    for iteration in cycle_summary.iterations[:-1]]))
        for cycle_summary, ok in cycle_summary_ok_list]

    cycle_summary_ok_max_distance_list.sort(
        key=lambda (cycle_summary, ok, max_distance): (
            not ok, max_distance))
    
    result, m = [], 0
    for (i, (cycle_summary, ok, max_distance)) in enumerate(cycle_summary_ok_max_distance_list):
        good = cycle_summary.path in known_good
        m += good
        result.append((cycle_summary.path, ok, max_distance, good, m, m / (i + 1)))

    if args.distance_ordered_csv:
        with open(args.distance_ordered_csv, "wb") as out_file:
            writer = csv.writer(out_file)
            writer.writerows(result)
        logger.info("Wrote %d in: %s", len(result), args.distance_ordered_csv)

    if args.distance_ordered_tsv:
        with open(args.distance_ordered_tsv, "wb") as out_file:
            writer = csv.writer(out_file, dialect="excel-tab")
            writer.writerows(result)
        logger.info("Wrote %d in: %s", len(result), args.distance_ordered_tsv)


if __name__ == "__main__":
    main(sys.argv[1:])
